/**
 * 
 */
package com.jszw.bcp.kernal.filter;

import java.io.IOException;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;

import org.jboss.logging.MDC;
import org.springframework.stereotype.Component;

/**
 * 
 * @author zhengzw
 */
@Component
public class RequestIDFilter implements Filter {

	@Override
	public void init(FilterConfig filterConfig) throws ServletException {

	}

	@Override
	public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
			throws IOException, ServletException {

		HttpServletRequest req = (HttpServletRequest) request;
		MutableHttpServletRequest mutableRequest = new MutableHttpServletRequest(req);

		String uuid = UUID.randomUUID().toString();
		
		mutableRequest.putHeader("uuid", uuid);
		MDC.put("uuid", uuid);

		chain.doFilter(mutableRequest, response);
	}

	@Override
	public void destroy() {

	}

	/**
	 * request customHeaders
	 * @author zhengzw
	 *
	 */
	final class MutableHttpServletRequest extends HttpServletRequestWrapper {

		private final Map<String, String> customHeaders = new HashMap<String, String>();

		public MutableHttpServletRequest(HttpServletRequest request) {
			super(request);
		}

		public void putHeader(String name, String value) {
			this.customHeaders.put(name, value);
		}

		public String getHeader(String name) {
			// check the custom headers first
			String headerValue = customHeaders.get(name);

			if (headerValue != null) {
				return headerValue;
			}
			// else return from into the original wrapped object
			return ((HttpServletRequest) getRequest()).getHeader(name);
		}

		public Enumeration<String> getHeaderNames() {
			// create a set of the custom header names
			Set<String> set = new HashSet<String>(customHeaders.keySet());

			// now add the headers from the wrapped request object
			Enumeration<String> e = ((HttpServletRequest) getRequest()).getHeaderNames();
			while (e.hasMoreElements()) {
				// add the names of the request headers into the list
				String n = e.nextElement();
				set.add(n);
			}

			// create an enumeration from the set and return
			return Collections.enumeration(set);
		}
	}
}
